import numpy as np
import scipy as sp
from scipy.optimize import least_squares
from scipy.linalg import qr, svd, solve_triangular
import warnings

def high_precision_solve(V, L_rho, method='svd_truncated', **kwargs):
    """
    高精度求解线性系统 S @ a = b，其中 S = V.conj().T @ V, b = V.conj().T @ L_rho
    
    Parameters:
    -----------
    V : ndarray, shape (m, n)
        输入矩阵
    L_rho : ndarray, shape (m,)
        右端向量
    method : str
        求解方法选择：
        - 'svd_truncated': 截断SVD (推荐)
        - 'qr_pivoting': QR分解with列主元
        - 'iterative_refinement': 迭代精化
        - 'tikhonov_adaptive': 自适应Tikhonov正则化
        - 'multiple_precision': 多精度计算
    
    Returns:
    --------
    a : ndarray
        解向量
    """
    
    # 计算正规方程矩阵
    S = V.conj().T @ V
    b = V.conj().T @ L_rho
    
    # 输出基本信息
    rank_S = np.linalg.matrix_rank(S)
    augmented = np.hstack((S, b.reshape(-1, 1)))
    rank_augmented = np.linalg.matrix_rank(augmented)
    cond_S = np.linalg.cond(S)
    
    print(f"rank(S) = {rank_S}")
    print(f"rank([S | b]) = {rank_augmented}")
    print(f"Condition number of S: {cond_S:.4e}")
    
    if method == 'svd_truncated':
        result = svd_truncated_solve(S, b, **kwargs)
    elif method == 'qr_pivoting':
        result = qr_pivoting_solve(V, L_rho, **kwargs)
    elif method == 'iterative_refinement':
        result = iterative_refinement_solve(S, b, **kwargs)
    elif method == 'tikhonov_adaptive':
        result = tikhonov_adaptive_solve(S, b, **kwargs)
    elif method == 'multiple_precision':
        result = multiple_precision_solve(S, b, **kwargs)
    else:
        raise ValueError(f"Unknown method: {method}")
    
    # 根据问题类型决定是否取实部
    # 如果原始问题期望实数解，则取实部；否则保持复数
    if np.allclose(np.imag(result), 0, atol=1e-10):
        return np.real(result)
    else:
        print("Warning: Solution has significant imaginary components")
        print(f"Max imaginary part: {np.max(np.abs(np.imag(result))):.4e}")
        return result

def svd_truncated_solve(S, b, tolerance=1e-12, verbose=True):
    """
    使用截断SVD求解，这是处理病态矩阵最稳定的方法
    修复复数矩阵处理
    """
    print("\n=== SVD Truncated Solve ===")
    
    # SVD分解
    U, s, Vh = svd(S, full_matrices=False)
    
    # 找到有效的奇异值
    s_max = np.max(s)
    valid_indices = s > tolerance * s_max
    n_valid = np.sum(valid_indices)
    
    print(f"Total singular values: {len(s)}")
    print(f"Valid singular values (> {tolerance:.2e} * σ_max): {n_valid}")
    if n_valid > 0:
        print(f"Smallest valid singular value: {s[valid_indices][-1]:.4e}")
    print(f"Largest singular value: {s_max:.4e}")
    print(f"Condition number (before truncation): {s_max/s[-1]:.4e}")
    if n_valid > 1:
        print(f"Condition number (after truncation): {s_max/s[valid_indices][-1]:.4e}")
    
    # 截断SVD求解 - 正确处理复数
    s_inv = np.zeros_like(s)
    s_inv[valid_indices] = 1.0 / s[valid_indices]
    
    # a = V @ diag(1/s) @ U^H @ b (注意共轭转置)
    a = Vh.T.conj() @ (s_inv * (U.T.conj() @ b))
    
    # 计算残差
    residual = np.linalg.norm(S @ a - b)
    print(f"Residual norm for S @ a - b: {residual:.4e}")
    
    return a  # 保持复数形式

def qr_pivoting_solve(V, L_rho, tolerance=1e-12, verbose=True):
    """
    使用QR分解with列主元直接求解 V @ a = L_rho
    避免计算 S = V^H @ V，减少条件数恶化
    修复复数矩阵处理
    """
    print("\n=== QR Pivoting Solve ===")
    
    # QR分解with列主元 - 对V^H进行QR分解
    Q, R, P = qr(V.T.conj(), pivoting=True)  # 注意共轭转置
    
    # 找到有效的对角元素
    diag_R = np.abs(np.diag(R))
    max_diag = np.max(diag_R)
    valid_indices = diag_R > tolerance * max_diag
    n_valid = np.sum(valid_indices)
    
    print(f"Matrix rank from QR: {n_valid}")
    if n_valid > 0:
        print(f"Smallest valid diagonal element: {diag_R[valid_indices][-1]:.4e}")
    print(f"Largest diagonal element: {max_diag:.4e}")
    
    if n_valid == 0:
        print("Warning: No valid diagonal elements found!")
        return np.zeros(V.shape[1], dtype=V.dtype)
    
    # 求解 R[:n_valid, :n_valid] @ y = Q[:, :n_valid]^H @ L_rho
    Q_valid = Q[:, :n_valid]
    R_valid = R[:n_valid, :n_valid]
    
    y = solve_triangular(R_valid, Q_valid.T.conj() @ L_rho)  # 注意共轭转置
    
    # 构造完整解
    a_permuted = np.zeros(V.shape[1], dtype=V.dtype)
    a_permuted[:n_valid] = y
    
    # 恢复原始顺序
    a = np.zeros_like(a_permuted)
    a[P] = a_permuted
    
    # 计算残差
    residual = np.linalg.norm(V @ a - L_rho)
    print(f"Residual norm for V @ a - L_rho: {residual:.4e}")
    
    # 同时计算正规方程的残差用于比较
    S = V.conj().T @ V
    b = V.conj().T @ L_rho
    normal_residual = np.linalg.norm(S @ a - b)
    print(f"Residual norm for S @ a - b: {normal_residual:.4e}")
    
    return a  # 保持复数形式

def iterative_refinement_solve(S, b, max_iter=10, tolerance=1e-14, verbose=True):
    """
    迭代精化求解法
    """
    print("\n=== Iterative Refinement Solve ===")
    
    # 初始解（使用最小二乘）
    a = sp.linalg.lstsq(S, b, lapack_driver='gelsy')[0]
    
    print("Initial solution computed")
    
    for i in range(max_iter):
        # 计算残差
        r = b - S @ a
        residual_norm = np.linalg.norm(r)
        
        print(f"Iteration {i+1}: Residual norm = {residual_norm:.4e}")
        
        if residual_norm < tolerance:
            break
            
        # 求解修正方程 S @ da = r
        try:
            da = sp.linalg.lstsq(S, r, lapack_driver='gelsy')[0]
            a = a + da
        except:
            print("Iterative refinement failed, using current solution")
            break
    
    final_residual = np.linalg.norm(S @ a - b)
    print(f"Final residual norm for S @ a - b: {final_residual:.4e}")
    
    return np.real(a)

def tikhonov_adaptive_solve(S, b, verbose=True):
    """
    自适应Tikhonov正则化 - 修复复数处理
    """
    print("\n=== Adaptive Tikhonov Regularization ===")
    
    # 使用L-curve方法选择正则化参数
    U, s, Vh = svd(S, full_matrices=False)
    
    # 候选正则化参数
    lambda_values = np.logspace(-15, -1, 50)
    residual_norms = []
    solution_norms = []
    
    for lam in lambda_values:
        # 正则化求解 - 修正复数处理
        s_reg = s / (s**2 + lam**2)
        a_reg = Vh.T.conj() @ (s_reg * (U.T.conj() @ b))
        
        residual_norms.append(np.linalg.norm(S @ a_reg - b))
        solution_norms.append(np.linalg.norm(a_reg))
    
    # 选择L-curve拐点处的参数 - 增加数值稳定性检查
    curvatures = []
    for i in range(1, len(lambda_values)-1):
        # 添加数值稳定性检查
        if (residual_norms[i] <= 0 or solution_norms[i] <= 0 or 
            residual_norms[i-1] <= 0 or solution_norms[i-1] <= 0 or
            residual_norms[i+1] <= 0 or solution_norms[i+1] <= 0):
            curvatures.append(0)
            continue
            
        # 计算曲率（简化版本）
        try:
            dx1 = np.log(solution_norms[i]) - np.log(solution_norms[i-1])
            dy1 = np.log(residual_norms[i]) - np.log(residual_norms[i-1])
            dx2 = np.log(solution_norms[i+1]) - np.log(solution_norms[i])
            dy2 = np.log(residual_norms[i+1]) - np.log(residual_norms[i])
            
            denominator = (dx1**2 + dy1**2)**1.5
            if denominator < 1e-15:  # 避免除零
                curvatures.append(0)
            else:
                curvature = abs(dx1*dy2 - dx2*dy1) / denominator
                curvatures.append(curvature)
        except (ValueError, RuntimeWarning):
            curvatures.append(0)
    
    if len(curvatures) == 0 or max(curvatures) == 0:
        # 如果L-curve失败，使用GCV准则
        print("L-curve method failed, using GCV criterion")
        return tikhonov_gcv_solve(S, b, U, s, Vh)
    
    optimal_idx = np.argmax(curvatures) + 1
    optimal_lambda = lambda_values[optimal_idx]
    
    print(f"Optimal regularization parameter: {optimal_lambda:.4e}")
    
    # 使用最优参数求解
    s_optimal = s / (s**2 + optimal_lambda**2)
    a_optimal = Vh.T.conj() @ (s_optimal * (U.T.conj() @ b))
    
    residual = np.linalg.norm(S @ a_optimal - b)
    print(f"Residual norm for S @ a - b: {residual:.4e}")
    
    return a_optimal  # 保持复数结果，让调用者决定是否取实部

def multiple_precision_solve(S, b, verbose=True):
    """
    多精度计算（使用更高精度的数据类型）
    """
    print("\n=== Multiple Precision Solve ===")
    
    # 转换为更高精度
    S_hp = S.astype(np.complex256) if hasattr(np, 'complex256') else S.astype(np.complex128)
    b_hp = b.astype(np.complex256) if hasattr(np, 'complex256') else b.astype(np.complex128)
    
    print(f"Using precision: {S_hp.dtype}")
    
    # SVD求解
    try:
        U, s, Vh = svd(S_hp, full_matrices=False)
        
        # 截断小的奇异值
        tolerance = 1e-15
        s_max = np.max(s)
        valid_indices = s > tolerance * s_max
        
        s_inv = np.zeros_like(s)
        s_inv[valid_indices] = 1.0 / s[valid_indices]
        
        a_hp = Vh.T @ (s_inv * (U.T @ b_hp))
        
        # 转换回标准精度
        a = a_hp.astype(np.complex128)
        
        residual = np.linalg.norm(S @ a - b)
        print(f"Residual norm for S @ a - b: {residual:.4e}")
        
        return np.real(a)
        
    except Exception as e:
        print(f"High precision solve failed: {e}")
        return svd_truncated_solve(S, b, tolerance=1e-12)

def levenberg_marquardt_solve(V, L_rho, lambda_init=1e-3, max_iter=100, tolerance=1e-12):
    """
    Levenberg-Marquardt方法求解非线性最小二乘
    修复复数处理 - 将复数问题转换为实数问题
    """
    print("\n=== Levenberg-Marquardt Solve ===")
    
    # 将复数问题转换为实数问题
    # [Re(V)  -Im(V)] [Re(a)]   [Re(L_rho)]
    # [Im(V)   Re(V)] [Im(a)] = [Im(L_rho)]
    
    V_real = np.vstack([np.real(V), np.imag(V)])
    V_imag = np.vstack([-np.imag(V), np.real(V)])
    V_expanded = np.hstack([V_real, V_imag])
    
    L_rho_expanded = np.concatenate([np.real(L_rho), np.imag(L_rho)])
    
    def residual_func(a_expanded):
        return V_expanded @ a_expanded - L_rho_expanded
    
    # 初始猜测
    a0 = np.zeros(V_expanded.shape[1], dtype=np.float64)
    
    try:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", ComplexWarning)
            result = least_squares(residual_func, a0, method='lm', 
                                 ftol=tolerance, xtol=tolerance, max_nfev=max_iter)
        
        # 转换回复数形式
        n_vars = V.shape[1]
        a_real = result.x[:n_vars]
        a_imag = result.x[n_vars:]
        a_complex = a_real + 1j * a_imag
        
        residual = np.linalg.norm(V @ a_complex - L_rho)
        print(f"Residual norm for Levenberg-Marquardt: {residual:.4e}")
        
        return a_complex
        
    except Exception as e:
        print(f"Levenberg-Marquardt failed: {e}")
        return np.zeros(V.shape[1], dtype=complex)

# 添加GCV方法作为备选
def tikhonov_gcv_solve(S, b, U=None, s=None, Vh=None):
    """
    使用GCV (Generalized Cross Validation) 准则选择正则化参数
    """
    if U is None or s is None or Vh is None:
        U, s, Vh = svd(S, full_matrices=False)
    
    def gcv_function(alpha):
        # 计算GCV函数值
        filter_factors = s**2 / (s**2 + alpha**2)
        numerator = np.sum((1 - filter_factors)**2 * (U.T.conj() @ b)**2)
        denominator = (np.sum(1 - filter_factors))**2
        return numerator / denominator
    
    # 在对数空间搜索最优参数
    alpha_values = np.logspace(-15, -1, 100)
    gcv_values = [gcv_function(alpha) for alpha in alpha_values]
    
    optimal_alpha = alpha_values[np.argmin(gcv_values)]
    print(f"Optimal regularization parameter (GCV): {optimal_alpha:.4e}")
    
    # 使用最优参数求解
    s_optimal = s**2 / (s**2 + optimal_alpha**2)
    a_optimal = Vh.T.conj() @ (s_optimal * (U.T.conj() @ b))
    
    residual = np.linalg.norm(S @ a_optimal - b)
    print(f"Residual norm for S @ a - b (GCV): {residual:.4e}")
    
    return a_optimal

# 示例使用函数
def compare_all_methods(V, L_rho):
    """
    比较所有方法的性能
    """
    print("="*60)
    print("COMPARING ALL HIGH-PRECISION METHODS")
    print("="*60)
    
    methods = [
        ('svd_truncated', {}),
        ('qr_pivoting', {}),
        ('iterative_refinement', {}),
        ('tikhonov_adaptive', {}),
        ('multiple_precision', {})
    ]
    
    results = {}
    
    for method_name, kwargs in methods:
        try:
            print(f"\n{'='*20} {method_name.upper()} {'='*20}")
            results[method_name] = high_precision_solve(V, L_rho, method=method_name, **kwargs)
        except Exception as e:
            print(f"Method {method_name} failed: {e}")
            results[method_name] = None
    
    # 也测试 Levenberg-Marquardt
    try:
        print(f"\n{'='*20} LEVENBERG-MARQUARDT {'='*20}")
        results['levenberg_marquardt'] = levenberg_marquardt_solve(V, L_rho)
    except Exception as e:
        print(f"Levenberg-Marquardt failed: {e}")
        results['levenberg_marquardt'] = None
    
    return results

# 使用示例
if __name__ == "__main__":
    # 加载您的数据
    # S = np.load(path+'S_matrix.npy').astype(np.complex128)
    # b = np.load(path+'b.npy').astype(np.complex128)
    # V = np.load(path+'V.npy').astype(np.complex128)
    # L_rho = np.load(path+'L_rho.npy').astype(np.complex128)
    
    # 运行所有方法比较
    # results = compare_all_methods(V, L_rho)
    
    # 或者单独使用最推荐的方法
    # best_solution = high_precision_solve(V, L_rho, method='svd_truncated', tolerance=1e-15)
    pass